#pragma once

#include "../base/list.h"
#include "lock.h"
#include <vector>
#include <chrono>
#include <functional>
#include <thread>

OPEN_JLIB_NS


struct Scheduler {

protected:
  typedef std::chrono::microseconds duration_t;
  typedef std::function<void(int64_t)> mission_t;

  inline static std::chrono::microseconds currentTime () noexcept {
    using namespace std::chrono;
    return duration_cast<microseconds>(system_clock::now().time_since_epoch());
  }


  struct Mission {

    inline Mission (mission_t &&func, int64_t time, int64_t rate, int64_t steps, int64_t order) noexcept
    : _func(std::move(func))
    , _time(time)
    , _rate(rate)
    , _step(0)
    , _steps(steps)
    , _order(order) {
    }

    inline ~Mission () noexcept {
    }

    inline int64_t _delay () noexcept {
      return _time - currentTime().count();
    }

    inline int64_t _compare (const Mission* other) noexcept {
      if (this == other)
        return 0;
      auto diff = _time - other->_time;
      if (diff != 0)
        return diff;
      return _order - other->_order;
    }

    mission_t _func;
    int64_t _time;
    int64_t _rate;
    int64_t _step;
    int64_t _steps;
    int64_t _order;
    int64_t _index;
  };


  void _insert (size_t index, Mission* item) noexcept {
    for (; index > 0; ) {
      auto parentIndex = (index - 1) >> 1;
      auto parent = _queue[parentIndex];
      if (item->_compare(parent) >= 0)
        break;
      _queue[parent->_index = index] = parent;
      index = parentIndex;
    }
    _queue[item->_index = index] = item;
  }


  Mission* _poll () noexcept {
    auto head = _queue.front();
    auto half = _queue.size() >> 1;
    auto index = 0UL;
    for (; index < half; ) {
      auto childIndex = (index << 1) + 1;
      auto child = _queue[childIndex];
      auto nextChildIndex = childIndex + 1;
      if (nextChildIndex < _queue.size() && child->_compare(_queue[nextChildIndex]) > 0)
        child = _queue[childIndex = nextChildIndex];
      _queue[child->_index = index] = child;
      index = childIndex;
    }
    if (index < _queue.size() - 1) {
      auto tail = _queue.back();
      _insert(index, tail);
    }
    _queue.pop_back();
    return head;
  }

  void _push (Mission* item) noexcept {
    auto index = _queue.size();
    _queue.emplace_back(nullptr);
    _insert(index, item);
  }


  inline void _decIdleWorker () noexcept {
    std::lock_guard<SpinLock> lock{_status_lock};
    --_idle_worker_count;
  }


  inline void _incIdleWorker () noexcept {
    std::lock_guard<SpinLock> lock{_status_lock};
    ++_idle_worker_count;
  }

  struct Worker {

    inline Worker (Scheduler* scheduler, List *prev, List *next, bool core) noexcept
    : _thread(&Worker::run, this)
    , _scheduler(scheduler)
    , _list(prev, next)
    , _core(core) {
      _thread.detach();
    }

    inline ~Worker () noexcept {
    }

    void run () noexcept {
      _scheduler->_incIdleWorker();
      while (!_scheduler->_closing) {
        Mission* mission;
        if (_core) {
          mission = _scheduler->_takeMission(this);
          if (!mission)
            continue;
        }
        else {
          mission = _scheduler->_tryTakeMission(this);
          if (!mission)
            break;
        }
        _scheduler->_decIdleWorker();
        mission->_func(mission->_step++);
        _scheduler->_incIdleWorker();
        _scheduler->_postMission(mission);
      }
      _scheduler->_decIdleWorker();
      _scheduler->_removeWorker(this);
    }

    std::thread _thread;
    Scheduler* _scheduler;
    List _list;
    bool _core;
  };


  inline int _shouldNewWorker () noexcept {
    std::lock_guard<SpinLock> lock{_status_lock};
    if (_closing || _idle_worker_count > 0)
      return -1;
    ++_pool_size;
    return _pool_size <= _core_pool_size ? 1 : 0;
  }

  inline void _decWorkerCount () noexcept {
    std::lock_guard<SpinLock> lock{_status_lock};
    --_pool_size;
  }

  inline bool _shouldStartClose () noexcept {
    std::lock_guard<SpinLock> lock{_status_lock};
    if (_closed)
      return false;
    return _closing = true;
  }

  inline bool _checkAndClose () noexcept {
    std::lock_guard<SpinLock> lock{_status_lock};
    if (!_closed && _closing && !_pool_size)
      return _closed = true;
    return false;
  }

  bool _addWorker () noexcept {
    std::lock_guard<std::mutex> lock{_pool_mutex};
    auto hint = _shouldNewWorker();
    if (hint < 0) {
      _queue_condition.notify_one();
      return false;
    }
    new Worker(this, _pool.prev(), &_pool, (bool)hint);
    return true;
  }

  void _removeWorker (Worker* worker) noexcept {
    std::lock_guard<std::mutex> lock{_pool_mutex};
    worker->_list.detach();
    _decWorkerCount();
    delete worker;
    if (_checkAndClose())
      _pool_close_condition.notify_all();
  }


  void _addMission (mission_t &&func, int64_t time, int64_t rate, int64_t step, int64_t order) noexcept {
    std::unique_lock<std::mutex> lock{_queue_mutex};
    auto mission = new Mission(std::move(func), time, rate, step, order);
    _push(mission);
  }

  Mission* _takeMission (Worker* worker) noexcept {
    Mission* result = nullptr;
    {
      std::unique_lock<std::mutex> lock{_queue_mutex};
      while (!_closing) {
        if (_queue.empty()) {
          _queue_condition.wait(lock);
          continue;
        }
        auto first = _queue.front();
        auto delay = first->_delay();
        if (delay <= 0) {
          _leader = nullptr;
          result = _poll();
          break;
        }
        if (_leader && _leader != worker) {
          _queue_condition.wait(lock);
          continue;
        }
        _leader = worker;
        _queue_condition.wait_for(lock, duration_t(delay));
      }
    }
    if (result)
      _addWorker();
    return result;
  }

  Mission* _tryTakeMission (Worker* worker) noexcept {
    Mission* result = nullptr;
    {
      std::unique_lock<std::mutex> lock{_queue_mutex};
      while (!_closing) {
        if (_queue.empty()) {
          _queue_condition.wait_for(lock, _keep_alive_time);
          continue;
        }
        auto first = _queue.front();
        auto delay = first->_delay();
        if (delay <= 0) {
          _leader = nullptr;
          result = _poll();
          break;
        }
        if (_leader && _leader != worker) {
          _queue_condition.wait_for(lock, _keep_alive_time);
          continue;
        }
        _leader = worker;
        _queue_condition.wait_for(lock, duration_t(delay));
      }
    }
    if (result)
      _addWorker();
    return result;
  }

  void _postMission (Mission* mission) noexcept {
    if (--mission->_steps && mission->_rate) {
      mission->_time = mission->_rate > 0 ? mission->_time + mission->_rate : currentTime().count() - mission->_rate;
      {
        std::lock_guard<std::mutex> lock{_queue_mutex};
        _push(mission);
        if (_queue.front() == mission)
          _leader = nullptr;
      }
      _addWorker();
    }
    else delete mission;
  }


public:

  inline explicit Scheduler (size_t coreThreadSize = std::thread::hardware_concurrency(), int64_t keepAliveTime = 60000000) noexcept
  : _core_pool_size(coreThreadSize ? coreThreadSize : 1)
  , _keep_alive_time(keepAliveTime >= 0 ? keepAliveTime : 0) {
  }

  inline ~Scheduler () noexcept {
    shutdownNow();
  }


  inline void execute (mission_t &&func) noexcept {
    if (_closing)
      return;
    _addMission(std::forward<mission_t>(func), currentTime().count(), 0, 1, _order_seed++);
    _addWorker();
  }

  inline void execute (int64_t delay, mission_t &&func) noexcept {
    if (_closing)
      return;
    _addMission(std::forward<mission_t>(func), currentTime().count() + delay, 0, 1, _order_seed++);
    _addWorker();
  }

  inline void scheduleWithFixedRate (int64_t delay, int64_t rate, int64_t step, mission_t &&func) noexcept {
    if (_closing || step <= 0 || rate <= 0)
      return;
    _addMission(std::forward<mission_t>(func), currentTime().count() + delay, rate, step, _order_seed++);
    _addWorker();
  }

  inline void scheduleWithFixedDelay (int64_t delay, int64_t rate, int64_t step, mission_t &&func) noexcept {
    if (_closing || step <= 0 || rate <= 0)
      return;
    _addMission(std::forward<mission_t>(func), currentTime().count() + delay, -rate, step, _order_seed++);
    _addWorker();
  }


  inline void shutdown () noexcept {
    _closing = true;
  }

  void shutdownNow () noexcept {
    if (!_shouldStartClose())
      return;
    {
      std::unique_lock<std::mutex> lock {_queue_mutex};
      for (auto mission : _queue)
        delete mission;
      _queue.clear();
      _queue_condition.notify_all();
    }
    {
      std::unique_lock<std::mutex> lock {_pool_mutex};
      while (!_closed)
        _pool_close_condition.wait(lock);
    }
  }


private:
  std::vector<Mission*> _queue;
  std::condition_variable _queue_condition;
  std::mutex _queue_mutex;
  Worker* _leader {nullptr};

  List _pool;
  std::condition_variable _pool_close_condition;
  std::mutex _pool_mutex;
  size_t _pool_size {0};
  size_t _core_pool_size;
  size_t _idle_worker_count {0};

  SpinLock _status_lock;
  duration_t _keep_alive_time;
  std::atomic<int64_t> _order_seed {~0L};
  volatile bool _closing {false};
  volatile bool _closed {false};
};


CLOSE_JLIB_NS